import cv2
import numpy as np
import pyrealsense2 as rs
import torch

# 加载YOLOv5模型
model = torch.hub.load('ultralytics/yolov5', 'custom',
                       path='F:/Yolov5/yolov5-master/yolov5-master/runs/train/garbage_detection20/weights/best.pt')  # 使用你训练好的YOLOv5模型
model.eval()  # 设置为评估模式

# 创建RealSense管道
pipeline = rs.pipeline()
config = rs.config()

# 配置RealSense相机流：RGB和深度
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)

# 尝试启动管道
try:
    pipeline.start(config)
    print("RealSense相机已成功打开！")
except Exception as e:
    print(f"无法打开RealSense相机：{e}")
    exit(1)

# 获取相机的配置信息
device = pipeline.get_active_profile().get_device()
print(f"设备信息：{device.get_info(rs.camera_info.name)}")

# 进入循环获取实时帧并进行垃圾识别
try:
    while True:
        # 获取实时帧
        frames = pipeline.wait_for_frames()
        color_frame = frames.get_color_frame()
        depth_frame = frames.get_depth_frame()

        # 检查RGB图像是否有效
        if not color_frame:
            print("没有捕获到RGB图像！")
            continue

        # 将颜色图像转换为NumPy数组
        color_image = np.asanyarray(color_frame.get_data())

        # 使用YOLOv5进行垃圾检测
        results = model(color_image)  # 输入RGB图像
        # 获取YOLOv5的预测结果
        pred = results.pred[0]

        # 绘制边界框并显示结果
        for *box, conf, cls in pred:
            x1, y1, x2, y2 = map(int, box)  # 获取坐标
            label = model.names[int(cls)]  # 获取预测标签
            confidence = float(conf)  # 获取置信度

            # 绘制矩形框
            cv2.rectangle(color_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            # 显示标签和置信度
            cv2.putText(color_image, f'{label} {confidence:.2f}', (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

        # 显示图像
        cv2.imshow('RealSense RGB with Garbage Detection', color_image)

        # 按 'q' 键退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
finally:
    # 停止RealSense流
    pipeline.stop()
    cv2.destroyAllWindows()
